"""

Recursive k=2 clustering operation helper script


Copyright 2012 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import subprocess

import numpy, sys, os, getopt
from itertools import combinations as comb

try:
    import psyco
    psyco.full()
except:
    pass

import clustio, scripts

# Steps:
# Create output dir, change to output dir...remember that calling cluster will do this
# Cluster
# Call make_cluster_defs on either cores or consensus logs
# Cluster again using said defs, with -o number (can't call it cluster_# since that's already a file)

CORING = False
CORING_FILENAME = '2 cluster suggested cores.log'
NON_CORING_FILENAME = '2 clusters - 300 subsamples.log'

# FIXME: This should be like an ENV(cluster) to find the path
# Or we could call consensuscluster via import...
CLUSTER_PY_PATH = '/home/buttons/Projects/cluster/cluster.py'

STDERR = sys.stderr

def opts_override(opts, *args):
    # XXX This function banks on the fact that passing -f/--foo twice will result in the second one overriding the first

    o = list(opts)
    o.extend(args)
    return o

def run_cluster(opts, outdir, max_cluster_size):
    # Recursively calls itself to run clustering iterations

    ecode = subprocess.call(['python', CLUSTER_PY_PATH] + opts_override(opts, '-o', outdir, '-k', '2,2'))

    # FIXME Sometimes --coring will not return two clusters. Currently ignoring errors to work around
    # FIXME these cases. If --coring is not set, errors will still block recursion.
    if ecode and CORING:
        return

    if ecode:
        raise ValueError, 'Error occurred in clustering, aborting recursion.'

    result_dir  = os.path.realpath(outdir)
    current_dir = os.path.realpath(os.path.curdir)
    os.chdir(result_dir)

    # Call make_cluster_defs
    if CORING:
        fn = CORING_FILENAME
    else:
        fn = NON_CORING_FILENAME
    
    cd = clustio.read_cluster_log(fn)               # Format: "%d (color):"
    clusts = dict([ (x[:-1], cd[x]) for x in cd ])  # Remove trailing : in cluster names

    if len(clusts) == 2:
        # Don't bother recursing or making cluster definitions if the number of clusters isn't two
        # *Should* fix clustering errors when --coring is set
        # Handy for making an automatic stopping condition by setting --coring 1 and --max_cluster_size 1
        
        scripts.make_def_clusters_from_log(fn)
        
        for x in clusts:
            if len(clusts[x]) > max_cluster_size:
                color  = x.split()[1].replace('(', '').replace(')', '')
                number = x.split()[0]
    
                dirname = '%s_%s' % (number, color) # XXX Oct 23 2012 Format "%d (color)" for dirnames was causing problems, switch to "%d_color"
                
                # When this if statement is called twice (both clusters are worth breaking down),
                # we have to get back to the original outdir to follow that clustering tree
                os.chdir(result_dir)
    
                # make_def_clusters makes cluster_%d files, we convert from %d (color) output dir name
                run_cluster(opts_override(opts, '-c', 'cluster_%s' % number), dirname, max_cluster_size)
                
                # These stack after each run to allow recursing down the tree after completion
                os.chdir(current_dir)


if __name__ == '__main__':

    if not os.path.exists(CLUSTER_PY_PATH):
        try:
            os.path.exists(cluster.py)
            CLUSTER_PY_PATH = os.path.realpath('.')
        except:
            raise ValueError, 'Unable to find cluster.py! Edit path variable in recursive_cluster.py or run from cluster base folder.'

    opts = sys.argv[1:]
    
    outdir = None
    coring = False
    max_cluster_size = 100
    new_opts = list(opts)

    # Find the output dir, coring options, and cluster size threshold
    for i in xrange(len(opts)):
        if opts[i] == '-o':
            outdir = opts[i+1]
        elif opts[i] == '--coring':
            CORING = True
        elif opts[i] == '--max_cluster_size':
            max_cluster_size = int(opts[i+1])
        elif opts[i] == '--subsamples':
            NON_CORING_FILENAME = '2 clusters - %s subsamples.log' % opts[i+1]
        elif opts[i] == '--flist':
            new_opts = opts_override(new_opts, '--flist', os.path.realpath(opts[i+1]))
        elif opts[i] == '-f':
            new_opts = opts_override(new_opts, '-f', ','.join([ os.path.realpath(x) for x in opts[i+1].split(',') if x ]))
        elif opts[i] == '-c':
            new_opts = opts_override(new_opts, '-c', ','.join([ os.path.realpath(x) for x in opts[i+1].split(',') if x ]))

    if outdir is None:
        print('-o option required. See python cluster.py --help for usage')
        sys.exit(2)

    run_cluster(new_opts, outdir, max_cluster_size)
